package clw.toy;

import clw.Toy;

import java.util.*;

/**
 * @author chengliwei
 */
public class Resolver implements Expr.Visitor<Void>, Stmt.Visitor<Void> {
    private final Interpreter interpreter;
    private final Stack<Map<String, Boolean>> scopes = new Stack<>();
    private FunctionType currentFunction = FunctionType.NONE;
    private ClassType currentClass = ClassType.NONE;

    public Resolver(Interpreter interpreter) {
        this.interpreter = interpreter;
    }

    private enum FunctionType {
        NONE,
        METHOD,
        FUNCTION,
        INITIALIZER
    }

    private enum ClassType {
        NONE,
        CLASS,
        SUBCLASS
    }

    @Override
    public Void visitBinary(Expr.Binary binary) {
        resolve(binary.left);
        resolve(binary.right);
        return null;
    }

    @Override
    public Void visitGrouping(Expr.Grouping grouping) {
        resolve(grouping.expression);
        return null;
    }

    @Override
    public Void visitLiteral(Expr.Literal literal) {
        return null;
    }

    @Override
    public Void visitUnary(Expr.Unary unary) {
        resolve(unary.right);
        return null;
    }

    @Override
    public Void visitVariableExpr(Expr.Variable variable) {
        if (!scopes.isEmpty() && scopes.peek().get(variable.name.lexeme) == Boolean.FALSE) {
            Toy.error(variable.name, "Can't read local variable in its own initializer.");
        }
        resolveLocal(variable, variable.name);
        return null;
    }

    private void resolveLocal(Expr expr, Token name) {
        for (int i = scopes.size() - 1; i >= 0; i--) {
            if (scopes.get(i).containsKey(name.lexeme)) {
                interpreter.resolve(expr, scopes.size() - 1 - i);
                return;
            }
        }
    }

    @Override
    public Void visitAssignExpr(Expr.Assign assign) {
        resolve(assign.value);
        resolveLocal(assign, assign.name);
        return null;
    }

    @Override
    public Void visitThisExpr(Expr.This thisExpr) {
        if (currentClass == ClassType.NONE) {
            Toy.error(thisExpr.keyword, "Can't use 'this' outside of a class.");
        }
        resolveLocal(thisExpr, thisExpr.keyword);
        return null;
    }

    @Override
    public Void visitLogicalExpr(Expr.Logical logical) {
        resolve(logical.left);
        resolve(logical.right);
        return null;
    }

    @Override
    public Void visitCallExpr(Expr.Call call) {
        resolve(call.callee);
        for (Expr argument : call.arguments) {
            resolve(argument);
        }
        return null;
    }

    @Override
    public Void visitGetExpr(Expr.Get get) {
        resolve(get.object);
        return null;
    }

    @Override
    public Void visitSetExpr(Expr.Set set) {
        resolve(set.value);
        resolve(set.object);
        return null;
    }

    @Override
    public Void visitSuperExpr(Expr.Super superExpr) {
        if (currentClass == ClassType.NONE) {
            Toy.error(superExpr.keyword, "Can't use 'super' outside of a class.");
        } else if (currentClass != ClassType.SUBCLASS) {
            Toy.error(superExpr.keyword, "Can't use 'super' in a class with no superclass.");
        }
        resolveLocal(superExpr, superExpr.keyword);
        return null;
    }

    @Override
    public Void visitSelfIncrementOrDecrement(Expr.SelfIncrementOrDecrement expr) {
        resolve(expr.left);
        return null;
    }

    @Override
    public Void visitExpressionStmt(Stmt.Expression expression) {
        resolve(expression.expr);
        return null;
    }

    @Override
    public Void visitPrintStmt(Stmt.Print print) {
        resolve(print.value);
        return null;
    }

    @Override
    public Void visitVarStmt(Stmt.Var var) {
        declare(var.name);
        if (var.initializer != null) {
            resolve(var.initializer);
        }
        define(var.name);
        return null;
    }

    private void declare(Token name) {
        if (scopes.isEmpty()) return;
        Map<String, Boolean> scope = scopes.peek();
        if (scope.containsKey(name.lexeme)) {
            Toy.error(name, "Already a variable with this name in this scope.");
        }
        scope.put(name.lexeme, false);
    }

    private void define(Token name) {
        if (scopes.isEmpty()) return;
        scopes.peek().put(name.lexeme, true);
    }

    @Override
    public Void visitBlockStmt(Stmt.Block block) {
        beginScope();
        resolve(block.statements);
        endScope();
        return null;
    }

    public void resolve(List<Stmt> statements) {
        for (Stmt statement : statements) {
            resolve(statement);
        }
    }

    private void resolve(Stmt stmt) {
        stmt.accept(this);
    }

    private void resolve(Expr expr) {
        expr.accept(this);
    }

    private void beginScope() {
        scopes.push(new HashMap<>());
    }

    private void endScope() {
        scopes.pop();
    }

    @Override
    public Void visitClassStmt(Stmt.Class clazz) {
        ClassType enclosingClass = currentClass;
        currentClass = ClassType.CLASS;

        declare(clazz.name);
        define(clazz.name);

        if (clazz.superClass != null && clazz.name.lexeme.equals(clazz.superClass.name.lexeme)) {
            Toy.error(clazz.superClass.name, "A class can't inherit from itself");
        }

        if (clazz.superClass != null) {
            currentClass = ClassType.SUBCLASS;
            resolve(clazz.superClass);
            beginScope();
            scopes.peek().put("super", true);
        }

        beginScope();
        scopes.peek().put("this", true);

        for (Stmt.Function method : clazz.methods) {
            FunctionType type = FunctionType.METHOD;
            if (method.name.lexeme.equals("init")) {
                type = FunctionType.INITIALIZER;
            }
            resolveFunction(method, type);
        }
        endScope();

        if (clazz.superClass != null) {
            endScope();
        }

        currentClass = enclosingClass;
        return null;
    }

    @Override
    public Void visitFunctionStmt(Stmt.Function function) {
        declare(function.name);
        define(function.name);
        resolveFunction(function, FunctionType.FUNCTION);
        return null;
    }

    private void resolveFunction(Stmt.Function function, FunctionType type) {
        FunctionType enclosingFunction = currentFunction;
        currentFunction = type;

        beginScope();
        for (Token param : function.params) {
            declare(param);
            define(param);
        }
        resolve(function.body);
        endScope();

        currentFunction = enclosingFunction;
    }

    @Override
    public Void visitIfStmt(Stmt.If ifStmt) {
        resolve(ifStmt.condition);
        resolve(ifStmt.thenBranch);
        if (ifStmt.elseBranch != null) resolve(ifStmt.elseBranch);
        return null;
    }

    @Override
    public Void visitReturnStmt(Stmt.Return returnStmt) {
        if (currentFunction == FunctionType.NONE) {
            Toy.error(returnStmt.keyword, "Can't return from top-level code.");
        }
        if (returnStmt.value != null) {
            if (currentFunction == FunctionType.INITIALIZER) {
                Toy.error(returnStmt.keyword, "Can't return a value from an initializer.");
            }
            resolve(returnStmt.value);
        }
        return null;
    }

    @Override
    public Void visitWhileStmt(Stmt.While whileStmt) {
        resolve(whileStmt.condition);
        resolve(whileStmt.body);
        return null;
    }

}
